import socket
import json
import time
import copy

ROBOT_IP = "192.168.1.233"
P_REF = [0, -90, 90, -90, 90, 0]


class ETController:
    def __init__(self, ip, port=8055):
        self.ip = ip
        self.port = port
        self.sock = None

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.sock.connect((self.ip, self.port))
            return True
        except Exception as e:
            print(f"Failed to connect: {e}")
            self.sock.close()
            self.sock = None
            return False

    def disconnect(self):
        if self.sock:
            self.sock.close()
            self.sock = None

    def send_cmd(self, cmd, params=None, id=1):
        if not params:
            params = []
        else:
            params = json.dumps(params)
        send_str = (
            '{{"method":"{0}","params":{1},"jsonrpc":"2.0","id":{2}}}'.format(
                cmd, params, id
            )
            + "\n"
        )
        try:
            self.sock.sendall(bytes(send_str, "utf-8"))
            ret = self.sock.recv(1024)
            jdata = json.loads(str(ret, "utf-8"))
            if "result" in jdata.keys():
                return True, json.loads(jdata["result"]), jdata["id"]
            elif "error" in jdata.keys():
                return False, jdata["error"], jdata["id"]
            else:
                return False, None, None
        except Exception as e:
            print(f"Failed to send command: {e}")
            return False, None, None

    def clear_alarm(self):
        ret, result, id = self.send_cmd("clearAlarm")
        if ret:
            print("Cleared Alarm")
        else:
            print("Failed to clear alarm")
        return ret, result, id

    def get_motor_status(self):
        ret, result, id = self.send_cmd("getMotorStatus")
        if ret:
            print("Got Motor Status")
            print("Result:", result)
        else:
            print("Failed to get motor status")
        return ret, result, id

    def sync_motor_status(self):
        ret, result, id = self.send_cmd("syncMotorStatus")
        if ret:
            print("Synced Motor Status")
        else:
            print("Failed to sync motor status")
        return ret, result, id

    def set_servo_status(self, status):
        ret, result, id = self.send_cmd("set_servo_status", {"status": status})
        if ret:
            print("Set Servo Status")
            print("Result:", result)
        else:
            print("Failed to set servo status")
        return ret, result, id

    def get_servo_status(self):
        ret, result, id = self.send_cmd("getServoStatus")
        if ret:
            print("Got Servo Status")
            print("Result:", result)
        else:
            print("Failed to get servo status")
        return ret, result, id

    def get_tcp_pose(self):
        success_pose, hand_pose, cmd_id_pose = self.send_cmd("get_tcp_pose",{"coordinate_num": -1, "tool_num": 0, "unit_type": 0})
        x = hand_pose[0]
        y = hand_pose[1]
        z = hand_pose[2]
        rx = hand_pose[3]
        ry = hand_pose[4]
        rz = hand_pose[5]
        return success_pose, (x,y,z,rx,ry,rz), cmd_id_pose
    
    def robot_init(self):
        # Clear Alarm
        self.clear_alarm()
        time.sleep(1)

        # Get Motor Status
        ret, result, _ = self.get_motor_status()
        if ret and result != 1:
            # Sync Motor Status
            self.sync_motor_status()
            time.sleep(0.5)
        # Set Servo Status
        self.set_servo_status(1)
        time.sleep(1)
        # Get Servo Status
        self.get_servo_status()

    def move_relative_deviation(self, last_pose, relative_deviation):
        if (
            last_pose is None
            or relative_deviation is None
            or len(last_pose) != 6
            or len(relative_deviation) != 6
        ):
            print("Move relative deviation input invalid !!!")
            return False
        curr_pose = []
        for numrate, devi in enumerate(relative_deviation):
            curr_pose.append(last_pose[numrate] + devi)
        suc_inverse, result_inverse, id = self.send_cmd(
            "inverseKinematic",
            {
                "targetPose": curr_pose,
                "referencePos": P_REF,
                "unit_type": 0,
            },
        )
        if suc_inverse:
            print("Inverse kinematic result is:", result_inverse)
            suc_move, result_move, id = self.send_cmd(
                "moveByLine",
                {
                    "targetPos": result_inverse,
                    "speed_type": 0,
                    "speed": 20,
                },
            )
            time.sleep(5)
            if suc_move:
                print("Current pose is:", curr_pose)
                # 获取机器人当前位姿
                # 帅：如果机械臂正在移动，get_tcp_pose获取的是移动前的位置，可以添加延迟，等机械臂停了，再get_tcp_pose获取
                _,curr_pose_get,_ = self.get_tcp_pose()
                print("Get current pose :", curr_pose_get)
                return True, curr_pose
            else:
                print("Failed to move to target pose")
                return False, last_pose
        else:
            print("Failed to get inverse kinematic solution")
            return False, last_pose


if __name__ == "__main__":
    try:
        # Example usage
        controller = ETController(ROBOT_IP)
        if controller.connect():
            controller.robot_init()
            # 获取机器人状态
            success, result, cmd_id = controller.send_cmd("getRobotState")
            # 打印结果
            if success:
                # 获取机器人当前位姿
                success_pose,hand_pose,_ = controller.get_tcp_pose()
                print("hand pose result:", hand_pose)
                if success_pose:
                    last_pose = copy.deepcopy(hand_pose)
                    while True:
                        try:
                            relative_deviation = []
                            # 提示用户输入位置偏移量
                            print("Enter position deviation:")
                            for i in range(3):
                                value = float(input(f"Axis {i + 1}: "))
                                relative_deviation.append(value)
                            # 提示用户输入角度偏移量
                            print("Enter angle deviation:")
                            for i in range(3):
                                value = float(input(f"Axis {i + 1}: "))
                                relative_deviation.append(value)

                            suc_move, curr_pose_get = controller.move_relative_deviation(
                                    last_pose, relative_deviation
                                )
                            
                            if suc_move:
                                print(
                                    "SUCCESS !!! Please input new position deviation."
                                )
                            else:
                                print("ERROR !!! Please try again.")
                            last_pose = copy.deepcopy(curr_pose_get)
                        except ValueError:
                            print("Please enter valid numbers.")
                else:
                    print("Failed to get TCP pose")
            else:
                print("Failed to get robot state")
        else:
            print("Failed to connect to ET Controller")
    except KeyboardInterrupt:
        pass
    finally:
        controller.disconnect()
